{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from langchain.prompts.example_selector.base import BaseExampleSelector\n",
    "from typing import Dict, List\n",
    "import numpy as np\n",
    "\n",
    "# 定义一个自定义示例选择器\n",
    "class CustomExampleSelector(BaseExampleSelector):\n",
    "    def __init__(self, examples: List[Dict[str, str]]):\n",
    "        self.examples = examples\n",
    "\n",
    "    def add_example(self, example: Dict[str, str]) -> None:\n",
    "        \"\"\"添加新的示例到examples列表中\"\"\"\n",
    "        self.examples.append(example)\n",
    "\n",
    "    def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:\n",
    "        \"\"\"从examples中随机选择指定数量的示例\"\"\"\n",
    "        return np.random.choice(self.examples, size=2, replace=False)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 示例列表\n",
    "examples = [\n",
    "    {\"input\": \"happy\", \"output\": \"sad\"},\n",
    "    {\"input\": \"tall\", \"output\": \"short\"},\n",
    "    {\"input\": \"energetic\", \"output\": \"lethargic\"},\n",
    "    {\"input\": \"sunny\", \"output\": \"gloomy\"},\n",
    "]\n",
    "\n",
    "# 实例化 CustomExampleSelector\n",
    "example_selector = CustomExampleSelector(examples)\n",
    "\n",
    "# 添加新的示例\n",
    "example_selector.add_example({\"input\": \"windy\", \"output\": \"calm\"})\n",
    "print(example_selector.examples)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import FewShotPromptTemplate\n",
    "from langchain import PromptTemplate\n",
    "\n",
    "# 定义模板\n",
    "example_formatter_template = \"\"\"input: {input} output: {output}\\\\n\"\"\"\n",
    "\n",
    "example_prompt = PromptTemplate(\n",
    "    input_variables=[\"input\", \"output\"],\n",
    "    template=example_formatter_template,\n",
    ")\n",
    "\n",
    "# 创建 FewShotPromptTemplate\n",
    "similar_prompt = FewShotPromptTemplate(\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=\"Give the antonym of every input\",\n",
    "    suffix=\"Input: {input}\\\\nOutput:\",\n",
    "    input_variables=[\"input\"],\n",
    ")\n",
    "\n",
    "# 格式化提示\n",
    "formatted_prompt = similar_prompt.format(input=\"beautiful\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.chat_models import ChatTongyi\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "# 创建聊天模型实例\n",
    "# 加载环境变量\n",
    "load_dotenv()\n",
    "# 初始化聊天模型\n",
    "chat = ChatTongyi(\n",
    "    model='qwen-plus',\n",
    "    top_p=0.9,\n",
    "    temperature=0.9,\n",
    "    api_key=os.getenv(\"DASHSCOPE_API_KEY\")\n",
    ")\n",
    "response = chat.invoke(formatted_prompt)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入必要的库\n",
    "from langchain.prompts.example_selector.base import BaseExampleSelector\n",
    "from typing import Dict, List\n",
    "import numpy as np\n",
    "import difflib\n",
    "from langchain.prompts import FewShotPromptTemplate\n",
    "from langchain import PromptTemplate\n",
    "from langchain_community.chat_models import ChatTongyi\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# 定义类别及其关键词\n",
    "CATEGORIES = {\n",
    "    \"手机\": [\"手机\", \"智能手机\", \"iPhone\", \"Android\", \"移动\"],\n",
    "    \"电脑\": [\"电脑\", \"笔记本\", \"台式机\", \"Mac\", \"Windows\"],\n",
    "    \"服装\": [\"服装\", \"衬衫\", \"裤子\", \"连衣裙\", \"服饰\"],\n",
    "}\n",
    "\n",
    "# 定义自定义示例选择器\n",
    "class CustomExampleSelector(BaseExampleSelector):\n",
    "    def __init__(self, examples: List[Dict[str, str]], categories: Dict[str, List[str]]):\n",
    "        \"\"\"\n",
    "        初始化示例选择器。\n",
    "\n",
    "        :param examples: 示例列表，每个示例应包含 'input', 'output' 和 'category' 字段。\n",
    "        :param categories: 类别定义，每个类别对应相关的关键词列表。\n",
    "        \"\"\"\n",
    "        self.examples = examples\n",
    "        self.categories = categories\n",
    "\n",
    "    def add_example(self, example: Dict[str, str]) -> None:\n",
    "        \"\"\"\n",
    "        添加新的示例到examples列表中。\n",
    "\n",
    "        :param example: 包含 'input', 'output' 和 'category' 的字典。\n",
    "        \"\"\"\n",
    "        self.examples.append(example)\n",
    "\n",
    "    def categorize_input(self, input_text: str) -> str:\n",
    "        \"\"\"\n",
    "        根据输入文本识别其类别。\n",
    "\n",
    "        :param input_text: 用户输入的文本。\n",
    "        :return: 识别的类别名称，如果无法识别则返回 'general'。\n",
    "        \"\"\"\n",
    "        input_text_lower = input_text.lower()\n",
    "        for category, keywords in self.categories.items():\n",
    "            for keyword in keywords:\n",
    "                if keyword.lower() in input_text_lower:\n",
    "                    return category\n",
    "        return \"general\"\n",
    "\n",
    "    def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:\n",
    "        \"\"\"\n",
    "        根据输入变量选择相关的示例。\n",
    "\n",
    "        :param input_variables: 包含 'input' 的字典。\n",
    "        :return: 选定的示例列表。\n",
    "        \"\"\"\n",
    "        input_text = input_variables.get(\"input\", \"\")\n",
    "        category = self.categorize_input(input_text)\n",
    "        print(f\"识别的类别: {category}\")\n",
    "\n",
    "        if category != \"general\":\n",
    "            # 从识别的类别中筛选示例\n",
    "            category_examples = [ex for ex in self.examples if ex.get(\"category\") == category]\n",
    "            if len(category_examples) >= 2:\n",
    "                selected_examples = np.random.choice(category_examples, size=2, replace=False)\n",
    "            else:\n",
    "                # 如果该类别的示例不足两个，则使用所有可用的示例\n",
    "                selected_examples = category_examples\n",
    "        else:\n",
    "            # 如果无法识别类别，则随机选择两个示例\n",
    "            selected_examples = np.random.choice(self.examples, size=2, replace=False)\n",
    "\n",
    "        # 将 NumPy 的 ndarray 转换为列表\n",
    "        return selected_examples.tolist()\n",
    "\n",
    "# 示例列表，包含类别标签\n",
    "examples = [\n",
    "    {\"input\": \"iPhone 最新款\", \"output\": \"功能更强大的 Iphone 手机\", \"category\": \"手机\"},\n",
    "    {\"input\": \"MacBook Pro\", \"output\": \"性价比更高的 苹果 笔记本\", \"category\": \"电脑\"},\n",
    "    {\"input\": \"时尚衬衫\", \"output\": \"舒适的休闲裤\", \"category\": \"服装\"},\n",
    "    {\"input\": \"华为 Mate 系列\", \"output\": \"遥遥领先的手机\", \"category\": \"手机\"},\n",
    "    {\"input\": \"ThinkPad 笔记本\", \"output\": \"轻薄便携的办公电脑 \", \"category\": \"电脑\"},\n",
    "    {\"input\": \"夏季连衣裙\", \"output\": \"经典款牛仔裤\", \"category\": \"服装\"},\n",
    "    {\"input\": \"华硕台式机\", \"output\": \"高性能的自组装电脑\", \"category\": \"电脑\"},\n",
    "    {\"input\": \"运动服饰\", \"output\": \"休闲时尚的外套\", \"category\": \"服装\"},\n",
    "    {\"input\": \"OnePlus 手机\", \"output\": \"高端的 Google Pixel\", \"category\": \"手机\"},\n",
    "    {\"input\": \"Apple MacBook 笔记本\", \"output\": \"高端的 Apple MacBook，无敌的设计\", \"category\": \"电脑\"},\n",
    "]\n",
    "\n",
    "# 实例化 CustomExampleSelector\n",
    "example_selector = CustomExampleSelector(examples, CATEGORIES)\n",
    "\n",
    "# 添加新的示例\n",
    "example_selector.add_example({\"input\": \"小米手机\", \"output\": \"性价比无敌 \", \"category\": \"手机\"})\n",
    "print(\"更新后的示例列表:\")\n",
    "for ex in example_selector.examples:\n",
    "    print(ex)\n",
    "\n",
    "# 定义示例格式模板\n",
    "example_formatter_template = \"\"\"输入: {input} 输出: {output}\\\\n\"\"\"\n",
    "\n",
    "example_prompt = PromptTemplate(\n",
    "    input_variables=[\"input\", \"output\"],\n",
    "    template=example_formatter_template,\n",
    ")\n",
    "\n",
    "# 创建 FewShotPromptTemplate\n",
    "similar_prompt = FewShotPromptTemplate(\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=\"根据以下示例，为输入生成一个吸引人的宣传词。\",\n",
    "    suffix=\"输入: {input}\\\\n输出:\",\n",
    "    input_variables=[\"input\"],\n",
    ")\n",
    "\n",
    "# 格式化提示\n",
    "formatted_prompt = similar_prompt.format(input=\"iPhone 14 Pro\")\n",
    "print(\"\\\\n格式化后的提示:\\\\n\", formatted_prompt)\n",
    "\n",
    "# 初始化聊天模型\n",
    "load_dotenv()\n",
    "\n",
    "chat = ChatTongyi(\n",
    "    model='qwen-plus',\n",
    "    top_p=0.9,\n",
    "    temperature=0.9,\n",
    "    api_key=os.getenv(\"DASHSCOPE_API_KEY\")\n",
    ")\n",
    "\n",
    "# 调用模型并输出结果\n",
    "response = chat.invoke(formatted_prompt)\n",
    "print(\"\\\\n模型响应:\\\\n\", response.content)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chatchat",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
